import ipaddress
import mmap
import os
from dataclasses import dataclass
from typing import Optional, Union

from pathlib import Path

# XDB 文件布局常量（固定）
HEADER_INFO_LEN = 256
VECTOR_INDEX_ROWS = 256
VECTOR_INDEX_COLS = 256
VECTOR_INDEX_SIZE = 8  # 每个向量索引项占 8 字节（startPtr:4 + endPtr:4）
SEGMENT_INDEX_SIZE = (
    14  # 每个段索引项占 14 字节（startIP:4 + endIP:4 + dataLen:2 + dataPtr:4）
)


class XdbError(Exception):
    """XDB 搜索相关的基础异常。"""

    pass


@dataclass(frozen=True)
class Region:
    """可选：解析 region 字符串的结构体。
    ip2region 的 region 格式通常为: country|region|province|city|isp
    """

    # 国家
    country: str
    # 地区
    region: str
    # 省份
    province: str
    # 城市
    city: str
    # 运营商
    isp: str

    @staticmethod
    def parse(region_str: str) -> "Region":
        parts = (region_str or "").split("|")
        parts += [""] * (5 - len(parts))
        return Region(*parts[:5])

    def as_dict(self, zero_to_empty: bool = True) -> dict[str, str]:
        d = {
            "country": self.country,
            "region": self.region,
            "province": self.province,
            "city": self.city,
            "isp": self.isp,
        }
        if zero_to_empty:
            for k, v in d.items():
                if v == "0":
                    d[k] = ""
        return d

    def __str__(self):
        # 国家、地区、省、市、运营商
        return f"{self.country}|{self.region}|{self.province}|{self.city}|{self.isp}"


class XdbSearcher:
    """
    XDB 搜索器
    - 支持三种数据读取模式：
        1) 全量内存: preload_content=True（最快，大内存）
        2) 仅向量索引: preload_vector_index=True（折中）
        3) 默认 mmap 文件（节省内存，适合单机多进程共享文件）

    用法：
    - with XdbSearcher(db_path) as searcher: ...
    - searcher = XdbSearcher(db_path); ...; searcher.close()
    """

    def __init__(
        self,
        db_path: Union[str, Path] = None,
        *,
        preload_vector_index: bool = False,
        preload_content: bool = False,
        encoding: str = "utf-8",
    ):
        if not db_path:
            # 引用上级文件夹中的 xdb 文件
            current_path = os.path.abspath(__file__)
            # 获取current_path的上级目录
            parent_path = os.path.dirname(current_path)
            db_path = os.path.join(parent_path, "data", "ip2region.xdb")
        self.db_path = Path(db_path)
        if not self.db_path.is_file():
            raise XdbError(f"XDB file not found: {self.db_path}")

        if preload_content and preload_vector_index:
            # 有了整库内容，向量索引可直接从 content 读，不必重复加载
            preload_vector_index = False

        self._encoding = encoding
        self._file = None
        self._mmap: Optional[mmap.mmap] = None
        self._content: Optional[bytes] = None
        self._vector_index: Optional[bytes] = None
        self._closed = True

        try:
            self._open_file()
            if preload_content:
                self._content = self._read_all()
            elif preload_vector_index:
                self._vector_index = self.load_vector_index_from_file(self.db_path)
        except Exception:
            self.close()
            raise

    @staticmethod
    def region_to_dict(region_str: str, zero_to_empty: bool = True) -> dict[str, str]:
        parts = (region_str or "").split("|")
        parts += [""] * (5 - len(parts))
        d = {
            "country": parts[0],
            "region": parts[1],
            "province": parts[2],
            "city": parts[3],
            "isp": parts[4],
        }
        if zero_to_empty:
            for k, v in d.items():
                if v == "0":
                    d[k] = ""
        return d

    # ---------- 公开的静态方法（工具/复用） ----------
    @staticmethod
    def load_vector_index_from_file(db_path: Union[str, Path]) -> bytes:
        """加载向量索引（仅 256*256*8 字节）"""
        p = Path(db_path)
        with p.open("rb") as f:
            f.seek(HEADER_INFO_LEN)
            length = VECTOR_INDEX_ROWS * VECTOR_INDEX_COLS * VECTOR_INDEX_SIZE
            return f.read(length)

    @staticmethod
    def load_content_from_file(db_path: Union[str, Path]) -> bytes:
        """加载全部内容到内存"""
        p = Path(db_path)
        return p.read_bytes()

    # ---------- 上下文管理 ----------
    def __enter__(self) -> "XdbSearcher":
        return self

    def __exit__(self, exc_type, exc, tb):
        self.close()

    # ---------- 公开 API ----------
    # 3) search：未命中返回 None；as_dict=True 则输出 dict（含 "0"->""）
    def search(self, ip, *, as_dict: bool = False):
        ip_int = self._ip_to_int(ip)
        region_str = self._search_by_ip_int(ip_int)
        if region_str == "":
            return None
        return self.region_to_dict(region_str) if as_dict else region_str

    # 4) 批量：同样返回 None（未命中）
    def search_many(self, ips, *, as_dict: bool = False):
        return [self.search(ip, as_dict=as_dict) for ip in ips]

    # 5) 便捷：直接要 dict，未命中返回 None
    def search_dict(self, ip):
        return self.search(ip, as_dict=True)

    def search_region(self, ip: Union[str, int, ipaddress.IPv4Address]) -> Region:
        """查询并解析为结构体"""
        return Region.parse(self.search(ip))

    @property
    def closed(self) -> bool:
        return self._closed

    def close(self) -> None:
        if self._mmap is not None:
            try:
                self._mmap.close()
            finally:
                self._mmap = None
        if self._file is not None:
            try:
                self._file.close()
            finally:
                self._file = None
        self._content = None
        self._vector_index = None
        self._closed = True

    # ---------- 内部实现 ----------
    def _open_file(self) -> None:
        self._file = self.db_path.open("rb")
        # 使用只读 mmap，便于多进程共享、减少内存占用
        self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ)
        self._closed = False

    def _read_all(self) -> bytes:
        if self._mmap is not None:
            return self._mmap[:]  # 拷贝
        assert self._file is not None
        self._file.seek(0)
        return self._file.read()

    def _ip_to_int(self, ip: Union[str, int, ipaddress.IPv4Address]) -> int:
        if isinstance(ip, int):
            if 0 <= ip <= 0xFFFFFFFF:
                return ip
            raise XdbError(f"Invalid IPv4 integer: {ip}")
        if isinstance(ip, ipaddress.IPv4Address):
            return int(ip)
        # str
        try:
            return int(ipaddress.IPv4Address(ip))
        except ipaddress.AddressValueError as e:
            raise XdbError(f"Invalid IPv4 string: {ip}") from e

    def _read_buffer(self, offset: int, length: int) -> memoryview:
        """从 content/mmap/文件 读取一段数据并返回 memoryview"""
        if self._content is not None:
            return memoryview(self._content)[offset : offset + length]
        if self._mmap is not None:
            return memoryview(self._mmap)[offset : offset + length]
        # 理论上不会走到这里；作为兜底
        assert self._file is not None
        self._file.seek(offset)
        return memoryview(self._file.read(length))

    @staticmethod
    def _u32_le(buf: memoryview | bytes, offset: int = 0) -> int:
        """小端 4 字节无符号整数"""
        return int.from_bytes(buf[offset : offset + 4], "little", signed=False)

    @staticmethod
    def _u16_le(buf: memoryview | bytes, offset: int = 0) -> int:
        """小端 2 字节无符号整数"""
        return int.from_bytes(buf[offset : offset + 2], "little", signed=False)

    def _vector_index_range(self, ip_int: int) -> tuple[int, int]:
        """基于向量索引定位 [sPtr, ePtr]"""
        il0 = (ip_int >> 24) & 0xFF
        il1 = (ip_int >> 16) & 0xFF
        idx = il0 * VECTOR_INDEX_COLS * VECTOR_INDEX_SIZE + il1 * VECTOR_INDEX_SIZE

        if self._vector_index is not None:
            s_ptr = self._u32_le(self._vector_index, idx)
            e_ptr = self._u32_le(self._vector_index, idx + 4)
            return s_ptr, e_ptr

        # 从 content 或 mmap 读取
        base = HEADER_INFO_LEN + idx
        buf = self._read_buffer(base, VECTOR_INDEX_SIZE)
        s_ptr = self._u32_le(buf, 0)
        e_ptr = self._u32_le(buf, 4)
        return s_ptr, e_ptr

    def _search_by_ip_int(self, ip_int: int) -> str:
        s_ptr, e_ptr = self._vector_index_range(ip_int)

        # 如果没有命中，直接返回空
        if e_ptr < s_ptr:
            return ""

        # 在段索引中二分查找
        left = 0
        right = (e_ptr - s_ptr) // SEGMENT_INDEX_SIZE
        data_len = -1
        data_ptr = -1

        while left <= right:
            mid = (left + right) >> 1
            p = s_ptr + mid * SEGMENT_INDEX_SIZE
            seg = self._read_buffer(p, SEGMENT_INDEX_SIZE)

            sip = self._u32_le(seg, 0)
            if ip_int < sip:
                right = mid - 1
                continue

            eip = self._u32_le(seg, 4)
            if ip_int > eip:
                left = mid + 1
                continue

            # 命中段
            data_len = self._u16_le(seg, 8)  # 修复原代码的端序错误
            data_ptr = self._u32_le(seg, 10)
            break

        if data_ptr < 0 or data_len <= 0:
            return ""

        region_bytes = self._read_buffer(data_ptr, data_len)
        return bytes(region_bytes).decode(self._encoding, errors="replace")


if __name__ == "__main__":
    from pathlib import Path

    # # 1) mmap（默认，省内存）
    # with XdbSearcher() as s:
    #     for ip in ["1.2.3.4", "192.168.1.1"]:
    #         print(ip, "->", s.search(ip))
    #
    # # 2) 仅预加载向量索引（加速查找，内存占用 ~512KB）
    # with XdbSearcher(preload_vector_index=True) as s:
    #     print(s.search("8.8.8.8"))
    #
    # # 3) 全量加载（最快，但占用 ~数据库大小）
    # s = XdbSearcher(preload_content=True)
    # print(s.search_region("114.114.114.114"))
    # s.close()

    with XdbSearcher("data/ip2region.xdb") as s:
        print(s.search("1.2.3.4", as_dict=True))  # 命中 => "国家|大区|省|市|ISP"
        print(s.search("1.2.3.4", as_dict=True))  # 命中 => dict，且 "0" 字段已变 ""
        print(s.search("203.0.113.1"))  # 未命中 => None
        print(s.search_dict("127.0.0.1"))  # 未命中 => None
